package ratelimit

import (
	"context"
	"gitee.com/xiao_hange/go-admin-pkg/pkg/ratelimit"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

//拦截器

type InterceptorBuilder struct {
	limiter ratelimit.Limiter
	key     string
}

func NewInterceptorBuilder(limiter ratelimit.Limiter, key string) *InterceptorBuilder {
	return &InterceptorBuilder{
		limiter: limiter,
		key:     key,
	}
}

// BuildServerInterceptor 服务端
func (i *InterceptorBuilder) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		limited, err := i.limiter.Limited(ctx, i.key)
		if err != nil {
			// 需要添加对应的log
			return nil, status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		if limited {
			return nil, status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		return handler(ctx, req)
	}
}

// BuildClientInterceptor 客户端
func (i *InterceptorBuilder) BuildClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		limited, err := i.limiter.Limited(ctx, i.key)
		if err != nil {
			// 需要添加对应的log
			return status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		if limited {
			return status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}
